/* Copyright (c) Huawei Technologies Co., Ltd. 2024-2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/

#ifndef METADEF_CXX_ASCEND_IR_DEF_H
#define METADEF_CXX_ASCEND_IR_DEF_H

#include <string>
#include <memory>
#include "attr_store.h"
#include "graph/compute_graph.h"
#include "symbolic.h"
#include "graph/node.h"
#include "graph/anchor.h"
#include "debug/ge_util.h"
#include "graph/utils/op_desc_utils.h"
#include "external/graph/operator.h"
#include "graph/utils/type_utils.h"
#include "graph/ascend_ir/ascend_ir_check.h"

namespace ge {
namespace {
constexpr int64_t kIdNone = -1;
}
struct SizeVar {
  using Type = enum : int32_t  {
    kSizeTypeVar = 0,
    kSizeTypeConst = 1,
  };
  int64_t id{};
  std::string name;
  int64_t const_value{};
  Type type;
  explicit SizeVar(ge::Expression expr_other) : expr(std::move(expr_other)) {}
  // TTODO 需要评审
  ge::Expression expr;
};
using SizeVarPtr = std::shared_ptr<SizeVar>;

struct Axis {
  using Type = enum : int32_t {
    kAxisTypeOriginal,
    kAxisTypeBlockOuter,  // outer axis after split by multicore
    kAxisTypeBlockInner,  // inner axis after split by multicore
    kAxisTypeTileOuter,   // outer axis after split by one core
    kAxisTypeTileInner,   // inner axis after split by one core
    kAxisTypeMerged,
    kAxisTypeInvalid
  };

  int64_t id{kIdNone};    // axis id
  std::string name;  // axis name
  Type type{kAxisTypeInvalid};
  bool bind_block{false};
  ge::Expression size;
  int32_t align{-1};
  std::vector<int64_t> from;
  int64_t split_pair_other_id{kIdNone};
  // 自动融合场景的默认值，手写场景可以做配置，供ATT使用
  bool allow_oversize_axis{false};
  bool allow_unaligned_tail{true};
};
using AxisPtr = std::shared_ptr<Axis>;
using AxisId = int64_t;
enum class TransType : int64_t {
  kSplit = 0,
  kMerge,
  kValid
};
struct OneTransInfo {
  TransType trans_type;
  std::vector<AxisPtr> src_axis;
  std::vector<AxisPtr> dst_axis;
};
using TransInfoRoadOfGraph = std::vector<OneTransInfo>;

enum class ComputeType : int32_t {
  kComputeLoad,
  kComputeStore,
  kComputeReduceStore,
  kComputeElewise,
  kComputeBroadcast,
  kComputeReduce,
  kComputeTranspose,
  kComputeInvalid,
};

enum class ComputeUnit : int32_t {
  kUnitNone,
  kUnitMTE1,
  kUnitMTE2,
  kUnitMTE3,
  kUnitScalar,
  kUnitVector,
  kUnitCube,
  kUnitInvalid,
};

enum class ApiType : int32_t {
  kAPITypeBuffer, // Workspace/Data/Constant/IndexExpr/Output
  kAPITypeCompute, // Load/Store/ReduceStore/Elewise/BroadCast/Reduce/Transpose
  kAPITypeInvalid,
};

struct ApiInfo {
  ApiType type = ApiType::kAPITypeInvalid;
  ComputeType compute_type = ComputeType::kComputeInvalid;
  ComputeUnit unit = ComputeUnit::kUnitInvalid;
};

struct SchedInfo {
  int64_t exec_order{kIdNone};
  int64_t sub_exec_order{kIdNone};
  std::vector<int64_t> axis;
  int64_t loop_axis{kIdNone};
};

class AscNodeAttr : public ge::AttrGroupsBase {
 public:
  std::string name;
  std::string type;
  SchedInfo sched;
  ApiInfo api;
  AscNodeAttr() = default;
  ~AscNodeAttr() override = default;
  static AscNodeAttr &GetAttrGroup(ge::Operator &op);
  std::unique_ptr<AttrGroupsBase> Clone() override;
};

class AscDataNodeAttr : public AscNodeAttr {
 public:
  ~AscDataNodeAttr() override = default;
  static AscDataNodeAttr &GetAttrGroup(ge::Operator &op);
  int64_t index{-1};
};

class AscGraphAttr : public ge::AttrGroupsBase {
 public:
  // TTODO 待确认正式方案
  int64_t tiling_key = -1;
  std::vector<AxisPtr> axis;
  // TTODO 待正式方案后删除
  TransInfoRoadOfGraph trans_info_road;
  std::vector<SizeVarPtr> size_vars;
  std::unique_ptr<AttrGroupsBase> Clone() override;
};

enum class AllocType : int32_t { kAllocTypeGlobal, kAllocTypeL1, kAllocTypeL2, kAllocTypeBuffer, kAllocTypeQueue };

enum class MemHardware : int32_t {
  kMemHardwareGM,
  kMemHardwareUB,
};

enum class Position : int32_t {
  kPositionGM,
  kPositionVecIn,
  kPositionVecOut,
  kPositionVecCalc,
};

struct MemAttr {
  int64_t tensor_id = kIdNone;
  AllocType alloc_type = AllocType::kAllocTypeGlobal;
  Position position = Position::kPositionGM;
  MemHardware hardware = MemHardware::kMemHardwareGM;
  std::vector<int64_t> buf_ids;
  // TTODO 待删除
  std::string name;
};

struct MemQueueAttr {
  int64_t id = kIdNone;
  int64_t depth{-1};
  int64_t buf_num{-1};
  // TTODO 待删除
  std::string name;
};

struct MemBufAttr {
  int64_t id = kIdNone;
  // TTODO 待删除
  std::string name;
};

struct MemOptAttr {
  int64_t reuse_id = kIdNone;
  int64_t ref_tensor = kIdNone;
  int64_t merge_scope = kIdNone;
};

struct KeptUbBuffer {
  Expression size;
  MemAttr mem;
};

class AscTensorDataType {
 public:
  operator ge::DataType() const {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(tensor_desc_);
    return tensor_desc_->GetDataType();
  };
  void operator=(const ge::DataType &other) {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(tensor_desc_);
    tensor_desc_->SetDataType(other);
  };
 private:
  friend struct AscNodeOutputs;
  GeTensorDesc *tensor_desc_;
};

class AscTensorAttr : public ge::AttrGroupsBase {
 public:
  AscTensorDataType dtype;
  std::vector<int64_t> axis;
  std::vector<ge::Expression> repeats;
  std::vector<ge::Expression> strides;
  std::vector<ge::Expression> offsets;
  std::vector<int64_t> vectorized_axis;
  std::vector<ge::Expression> vectorized_strides;
  MemAttr mem;
  MemQueueAttr que;
  MemBufAttr buf;
  MemOptAttr opt;
  static AscTensorAttr &GetTensorAttr(ge::Operator *op, const uint32_t index);
  static AscTensorAttr &GetTensorAttr(const OutDataAnchor &output);
  std::unique_ptr<AttrGroupsBase> Clone() override;
};
}  // namespace ge

#endif  // METADEF_CXX_ASCEND_IR_DEF_H
